import errno
import time
import io
import struct
import fcntl
import termios
import sys

class ModifiableStreamTerm(object):
  """Handles colorized terminal output with lines at the end that change.

  There are lines that scroll up the screen like normal output ("log" lines) and
  lines that are replaced constantly (starting at the "first" one).

  Attributes:
    lines: The lines that we currently have down at the bottom.
    line_count: How many lines at the bottom are currently on the screen.
    buff: Queued up log output.
    output: The file-like object we write to.
  """
  def __init__(self):
    self.lines = []
    self.line_count = 0
    self.buff = io.StringIO()
    self.output = sys.stdout
    try:
      h, w = struct.unpack('hhhh',
                           fcntl.ioctl(1, termios.TIOCGWINSZ, b'\000' * 8))[:2]
      self._is_tty = True
    except IOError as ioerr:
      if ioerr.errno == errno.ENOTTY or ioerr.errno == errno.EINVAL:
        self._is_tty = False
      else:
        raise ioerr

  @property
  def is_tty(self):
    """Returns if stdout is a tty."""
    return self._is_tty

  def log(self, data):
    """Adds a line that goes above the "first" line and stays there.

    Args:
      data: The string to write (no newline on the end).
    """
    self.reset(reset_lines = False)
    data = str(data)
    self.buff.write(data + "\n")

  def flush(self):
    """Flushes all pending output."""
    if self._is_tty:
      h, w = struct.unpack('hhhh',
                           fcntl.ioctl(1, termios.TIOCGWINSZ, b'\000' * 8))[:2]
      if self.line_count == 0:
        for line in self.lines:
          if self.line_count > 0:
            self.buff.write("\n")
          self.buff.write(str(line))
          self.line_count += abs(len(line) - 1) // w + 1
        self.line_count = max(0, self.line_count - 1)
    self.output.write(self.buff.getvalue())
    self.buff = io.StringIO()

  def reset(self, reset_lines=True):
    """Resets the output back up through the "first" line.

    Args:
      reset_lines: Whether to reset the lines we print (defaults to True).
    """
    if self.line_count > 0:
      self.buff.write("\33[" + str(self.line_count) + "F\33[J")
      self.line_count = 0
    if reset_lines:
      self.lines = []

  def __exit__(self, a, b, c):
    if self.line_count != 0:
      self.output.write('\n')

  def __enter__(self):
    return self

def TermColor(col):
  return lambda data: ColorLine(col, data)

class ColorLine(object):
  colorBlank = "\33[0m"
  def __init__(self, colorEnd, data, length = None):
    self.colorEnd = colorEnd
    self.data = colorEnd + data
    self.__len = (length if length != None else len(data))

  def __add__(self, other):
    if(isinstance(other, ColorLine)):
      return ColorLine(other.colorEnd, self.data + other.data,
        self.__len + other.__len)
    else:
      return ColorLine(ColorLine.colorBlank, str(self) + other,
        self.__len + len(other))


  def __len__(self):
    return self.__len

  def __str__(self):
    if self.colorEnd == ColorLine.colorBlank:
      return self.data
    return self.data + ColorLine.colorBlank

green = TermColor("\33[32m")
red = TermColor("\33[31m")
yellow = TermColor("\33[33m")

blank = TermColor(ColorLine.colorBlank)
